'''
Author: jojo
Date: 2021-04-23 03:07:36
LastEditors: jojo
LastEditTime: 2021-06-18 13:59:38
FilePath: /waifu/interpolation.py
倍帧
'''

import os
import json
from const import RIFE_PATH
import shutil
import threading
from tqdm import tqdm
import time
import signal

import sys
from utils import global_var_manager as gv
from utils import cut,concat,remove_path, remove_file,extract_audio,cont_audio,read_video_info,get_fps
# src_dir = './workspace/test/src'
# dst_dir = './workspace/test/out'

# filename = 'test.mp4'

preset = 9

@DeprecationWarning
def full_interpolation(src_dir,dst_dir,filename,gpu='0'):
    """全视频直接倍帧

    Args:
        src_dir (str): 源视频文件所在目录
        dst_dir (str): 输出目录
        filename (str): 文件名，带后缀
        gpu(str): 使用的gpu，默认为0
    """

    signal.signal(signal.SIGINT,interupt_handler)
    signal.signal(signal.SIGHUP, interupt_handler)
    signal.signal(signal.SIGTERM, interupt_handler)
    gv.set_value(name='interpolated',value=False)
    file_realname = filename.split('.')[0]
    src_file_path = os.path.join(src_dir,filename)
    dst_file_path = os.path.join(dst_dir,filename)

    input_frame_dir = os.path.join(src_dir,file_realname + '_input_frames')
    output_frame_dir = os.path.join(dst_dir,file_realname + '_output_frames')
    

    if not os.path.exists(input_frame_dir):
        os.mkdir(input_frame_dir)

    if not os.path.exists(output_frame_dir):
        os.mkdir(output_frame_dir)
    
    print("extracting audio...")
    audio_path = os.path.join(src_dir,file_realname + '.m4a')
    extract_audio(src_file_path,audio_path)
    print("extracted!")

    print("decoding frames...")
    decode_frames(src_file_path,input_frame_dir)
    print("decoded!")
    
    # use thread to monitor the process
    threads = []
    threads.append(threading.Thread(target=interpolate,args=(input_frame_dir,output_frame_dir,gpu)))
    threads.append(threading.Thread(target=process_monitor,args=(input_frame_dir,output_frame_dir,src_file_path)))
    print("interpolating....")
    # multi-thread processing
    for t in threads:
        t.start()
        
    for t in threads:
        t.join()
        
    
    print("interporated!")

    print("concating audio to new video...")
    res_dict = read_video_info(src_dir,filename)
    cont_audio(output_frame_dir,audio_path,dst_file_path,res_dict)
    print("audio concated!")

    print(f"{src_file_path} has been interpolated to {dst_file_path}")
    remove_cache_with_audio(input_frame_dir,output_frame_dir,audio_path)
    
def clip_interpolation(src_dir,dst_dir,filename,res_dict,gpu='0'):
    """处理分段视频的倍帧

    Args:
        src_dir (str): 源视频文件所在目录
        dst_dir (str): 输出目录
        filename (str): 文件名，带后缀
        res_dict(dict): 原视频属性
        gpu(str): 使用的gpu，默认为0
    """
    signal.signal(signal.SIGINT,interupt_handler)
    signal.signal(signal.SIGHUP, interupt_handler)
    signal.signal(signal.SIGTERM, interupt_handler)
    gv.set_value(name='interpolated',value=False)
    file_realname = filename.split('.')[0]
    src_file_path = os.path.join(src_dir,filename)
    dst_file_path = os.path.join(dst_dir,filename)

    input_frame_dir = os.path.join(src_dir,file_realname + '_input_frames')
    output_frame_dir = os.path.join(dst_dir,file_realname + '_output_frames')
    

    if not os.path.exists(input_frame_dir):
        os.mkdir(input_frame_dir)

    if not os.path.exists(output_frame_dir):
        os.mkdir(output_frame_dir)
        
    print("decoding frames...")
    decode_frames(src_file_path,input_frame_dir)
    print("decoded!")
    
    # use thread to monitor the process
    threads = []
    threads.append(threading.Thread(target=interpolate,args=(input_frame_dir,output_frame_dir,gpu)))
    threads.append(threading.Thread(target=process_monitor,args=(input_frame_dir,output_frame_dir,src_file_path)))
    print("interpolating....")
    # multi-thread processing
    for t in threads:
        t.start()
        
    for t in threads:
        t.join()
        
    print("interporated!")
    
    # res_dict = read_video_info(src_dir,filename)
    cont_frame2video(output_frame_dir,dst_file_path,res_dict)
    remove_cache(input_frame_dir,output_frame_dir)
    
def interpolation(src,dst,file,gpu,cut_time):
    # 切割视频
    clip_path_list,output_abs_root_path,clip_name_list = cut(video_path=os.path.join(src,file),
                                                             cut_time=cut_time) # 获得切割后所有分段的路径列表

    # 获得原始视频的信息
    res_dict = read_video_info(src,file)

    #  分离音频
    audio_path = extract_audio(src,file)

    # interpolated output path
    interpolated_dst_path = os.path.join(output_abs_root_path,'dst')
    if not os.path.exists(interpolated_dst_path):
        os.mkdir(interpolated_dst_path)
    
    # 每个视频单独处理
    
    # 需要将待处理的视频路径合并
    interpolated_clip_path_list = []
    
    for clip_name in clip_name_list:
        clip_interpolation(output_abs_root_path,interpolated_dst_path,clip_name,res_dict,gpu)
        # interpolation(output_abs_root_path,interpolated_dst_path,clip_name,gpu)
        interpolated_clip_path_list.append(os.path.join(interpolated_dst_path,clip_name))

    # 处理好之后再合并起来
    concat(interpolated_clip_path_list,file,interpolated_dst_path)
    # concat(interpolated_clip_path_list,file,dst)
    
    # 与音频合并起来
    cont_audio(src_path=interpolated_dst_path,
               filename=file,
               audio_path=audio_path,
               dst_file_path=dst)
    
    # 清除分割的视频
    remove_path(output_abs_root_path)
    remove_file(audio_path)



def remove_cache_with_audio(input_frame_dir,output_frame_dir,audio_path):
    
    print(f"cleaning {audio_path}")
    try:
        os.remove(audio_path)
    except Exception as e:
        print(f"{audio_path} does not exist!")
        
    print(f"cleaning {input_frame_dir}")
    try:
        shutil.rmtree(input_frame_dir)
    except Exception as e:
        print(f"{input_frame_dir} does not exist!")
        
    print(f"cleaning {output_frame_dir}")
    try:
        shutil.rmtree(output_frame_dir)
    except Exception as e:
        print(f"{output_frame_dir} does not exist!")
        
def remove_cache(input_frame_dir,output_frame_dir):
    
    print(f"cleaning {input_frame_dir}")
    try:
        shutil.rmtree(input_frame_dir)
    except Exception as e:
        print(f"{input_frame_dir} does not exist!")
        
    print(f"cleaning {output_frame_dir}")
    try:
        shutil.rmtree(output_frame_dir)
    except Exception as e:
        print(f"{output_frame_dir} does not exist!")



def decode_frames(src_file_path,input_frame_dir):
    # decode all frames
    # dec_fra_com = f'ffmpeg -i {src_file_path} {input_frame_dir}/frame_%08d.png'
    dec_fra_com = f'ffmpeg -y -hwaccel cuvid -i {src_file_path} {input_frame_dir}/frame_%08d.png'
    os.system(dec_fra_com)
    gv.set_value('input',input_frame_dir)
    

# interpolate 2x frame count
def interpolate(input_frame_dir,output_frame_dir,gpu):
    inter_com = f'{RIFE_PATH} -i {input_frame_dir} -o {output_frame_dir} -j 100:100:100 -g {gpu} -m rife-v2.4'
    gv.set_value('output',output_frame_dir)
    os.system(inter_com)
    gv.set_value('interpolated',True)
    
def interupt_handler(signum,frame):
    print("user has interupt the progress")
    gv.set_value('interpolated',True)
    input_path = gv.get_value('input')
    output_path = gv.get_value('output')
    audio_path = gv.get_value('audio')
    print("cleaning cache...")
    remove_cache(input_frame_dir=input_path,output_frame_dir=output_path,audio_path=audio_path)
    sys.exit()

    
def cont_frame2video(output_frame_dir,dst_file_path,res_dict):
    fps = get_fps(res_dict)
    pix_fmt = res_dict['streams'][0]['pix_fmt']
    double_fps = fps*2
    
    # how to determine the crf value: https://zhuanlan.zhihu.com/p/250590703
    # use cpu
    int_with_aud_com = f'ffmpeg -y -framerate {double_fps} -i {output_frame_dir}/%08d.png -c:a copy -crf 22 -c:v libx264 -pix_fmt {pix_fmt} {dst_file_path}'

    # use gpu
    # int_with_aud_com = f'ffmpeg -y -framerate {double_fps} -i {output_frame_dir}/%08d.png -c:a copy -c:v h264_nvenc -preset {preset} -pix_fmt {pix_fmt} {dst_file_path}' # use gpu
    
    os.system(int_with_aud_com)



def process_monitor(input_frame_dir,output_frame_dir,filepath):
    """monitor the process

    Args:
        input_frame_dir (str): the original video frames path
        output_frame_dir (str): the target video frames path
    """
    total_frame_amount = len(os.listdir(input_frame_dir))
    target_frame_amount = 2*total_frame_amount
    current_frame_amount = 0
    previous_frame_amount = 0
    time.sleep(1)
    with tqdm(total=target_frame_amount,desc =f'\033[0;37;42mPID: {os.getpid()}\033[0m | Interpolating file:{filepath} process') as progress_bar:
        # get the current processed frame amount
        while not current_frame_amount==target_frame_amount and gv.get_value('interpolated') == False:
            current_frame_amount = len(os.listdir(output_frame_dir))
            delta = current_frame_amount - previous_frame_amount
            previous_frame_amount = current_frame_amount
            progress_bar.update(delta)
            time.sleep(1)
            
    print("interpolation process done!")
            

if __name__=='__main__':
    src_dir = './workspace/src/p1/'
    filename = '02.mp4'
    dst_dir = './workspace/dst/'
    # filename = 'test.mp4'
    # interpolation(src_dir,dst_dir,filename)
    extract_audio(os.path.join(src_dir,filename),os.path.join(dst_dir,'02.m4a'))
    res_dict = read_video_info(src_dir,filename)
    cont_audio(output_frame_dir=os.path.join(dst_dir,'02_output_frames'),audio_path=os.path.join(dst_dir,'02.m4a'),
               dst_file_path=os.path.join(dst_dir,'02.mp4'),res_dict=res_dict)